from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import os

class AESCipher:
    def __init__(self, key=None):
        self.key = key or os.environ["ENCRYPTION_KEY"].encode()
        
    def encrypt(self, data: bytes) -> bytes:
        """AES-GCM加密"""
        iv = os.urandom(12)
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.GCM(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()
        ciphertext = encryptor.update(data) + encryptor.finalize()
        return iv + encryptor.tag + ciphertext
    
    def decrypt(self, data: bytes) -> bytes:
        """AES-GCM解密"""
        iv, tag, ciphertext = data[:12], data[12:28], data[28:]
        cipher = Cipher(
            algorithms.AES(self.key),
            modes.GCM(iv, tag),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        return decryptor.update(ciphertext) + decryptor.finalize()